package com.feidee.fd.sml.algorithm.component.validation.assertion

import com.feidee.fd.sml.algorithm.component.validation.ValidationParam
import com.feidee.fd.sml.algorithm.util.DataFrameUtils
import org.apache.spark.SparkException
import org.apache.spark.sql.DataFrame

/**
  * @Author songhaicheng
  * @Date 2019/5/14 15:42
  * @Description
  * @Reviewer
  */
case class OddColAssertion(
                            override val op: String,
                            override val threshold: Int,
                            col: String,
                            function: String) extends Assertion {

  def this() = this("=", 0, "col", "size")

  val availableFunctions = Array("size", "distinct(size)")

  override def assert(data: DataFrame, validationParam: ValidationParam): Boolean = {
    // 避免将全量数据拉到本地，有 OOM 风险！
    val real = function.trim.toLowerCase match {
      case "size" =>
        DataFrameUtils.getColumnValues(data, col).length
      case "distinct(size)" =>
        DataFrameUtils.getColumnDistinctValues(data, col).length
      case _ => throw new SparkException(s"暂只支持 [${availableFunctions.mkString(", ")}] 方法")
    }

    println(s"read $real")

    val pass = op.trim.toLowerCase match {
      case ">" =>
        real > threshold
      case ">=" =>
        real >= threshold
      case "<" =>
        real < threshold
      case "<=" =>
        real <= threshold
      case "=" =>
        real == threshold
      case "!=" =>
        real != threshold
      case _ => throw new SparkException(s"暂只支持 [>, >=, <, <=, =, !=] 关系符，但设置的是 $op")
    }
    if (pass) {
      logInfo(s"假设列 $col 的 $function $op $threshold（阈值）通过，实际值 $real")
    } else {
      logError(s"假设列 $col 的 $function $op $threshold（阈值）不通过，实际值 $real")
    }
    pass
  }
}
